package main

import (
	"bufio"
	"context"
	"fmt"
	"net"
	"sync"
	"time"

	"github.com/spf13/pflag"
	"github.com/spf13/viper"
	"go.uber.org/atomic"
)

func init() {
	pflag.IntP("connections", "c", 1, "connections(goroutines) count")
	pflag.DurationP("time", "m", 15*time.Second, "test time")
	pflag.StringP("address", "a", "localhost:9999", "server address")
}

func main() {
	pflag.Parse()
	viper.BindPFlags(pflag.CommandLine)

	var tcpAddr *net.TCPAddr
	addr := viper.GetString("address")
	tcpAddr, _ = net.ResolveTCPAddr("tcp", addr)

	tm := viper.GetDuration("time")
	ctx, cancel := context.WithTimeout(context.Background(), tm)
	defer cancel()

	var wg sync.WaitGroup
	var transactions atomic.Uint32
	connections := viper.GetInt("connections")
	for i := 0; i < connections; i++ {
		wg.Add(1)
		go runStressTest(ctx, &wg, &transactions, tcpAddr)
	}
	wg.Wait()

	fmt.Printf("TPS: %d / %f = %f (t/s)\n", transactions.Load(), tm.Seconds(),
		float64(transactions.Load())/tm.Seconds())
}

func runStressTest(ctx context.Context, wg *sync.WaitGroup, transactions *atomic.Uint32, tcpAddr *net.TCPAddr) {
	defer wg.Done()

	runs := runStressTest2(ctx, tcpAddr)
	transactions.Add(runs)
}

func runStressTest2(ctx context.Context, tcpAddr *net.TCPAddr) uint32 {
	conn, err := net.DialTCP("tcp", nil, tcpAddr)
	if err != nil {
		fmt.Println("Client connect error ! " + err.Error())
		return 0
	}
	defer conn.Close()

	// fmt.Println(conn.LocalAddr().String() + " : Client connected!")
	i := uint32(0)
	for {
		select {
		case <-ctx.Done():
			return i
		default:
			if err := request(conn); err != nil {
				fmt.Println(err)
				return i
			}
			i++
		}
	}
}

func request(conn *net.TCPConn) error {
	b := []byte("Say hello to Server... \n")
	conn.Write(b)

	reader := bufio.NewReader(conn)
	_, err := reader.ReadString('\n')
	return err
}
